import torch
from torch import nn
from torch.nn import Sequential, Conv2d, MaxPool2d, Linear, Flatten


class MyModule(nn.Module):

    def __init__(self) -> None:
        super().__init__()
        self.module = Sequential(
            Conv2d(3, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 32, 5, padding=2),
            MaxPool2d(2),
            Conv2d(32, 64, 5, padding=2),
            MaxPool2d(2),
            Flatten(),
            Linear(1024, 64),
            Linear(64, 10)
        )

    def forward(self, input):
        return self.module(input)


if __name__ == '__main__':
    module = MyModule()
    input = torch.ones((64, 3, 32, 32))
    output = module(input)
    print(output.shape)
